package com.github.kezhenxu94.data.sink;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hbase.HBaseConfiguration;
import org.apache.hadoop.hbase.TableName;
import org.apache.hadoop.hbase.client.Connection;
import org.apache.hadoop.hbase.client.ConnectionFactory;
import org.apache.hadoop.hbase.client.Put;
import org.apache.hadoop.hbase.client.Table;

import java.io.IOException;
import java.util.Objects;
import java.util.Properties;
import java.util.concurrent.*;
import java.util.function.Function;

/**
 * Created by kezhenxu on 1/6/17.
 *
 * @author kezhenxu (kezhenxu94@163.com)
 */
public class HBaseSink<T> implements Sink<T> {

    protected ExecutorService executor;
    protected Configuration configuration;
    protected Connection connection;
    protected Properties properties;
    protected BlockingQueue<T> queue;
    protected volatile boolean isRunning;
    protected String tableName;
    protected String columnFamily;
    protected Function<T, byte[]> keySupplier;
    protected Function<T, byte[]> valueSupplier;

    @Override
    public void setup() throws Exception {
        this.configuration = HBaseConfiguration.create();
        this.properties = new Properties();
        this.properties.load(getClass().getResourceAsStream("/hbase-sink.properties"));
        this.properties.entrySet().forEach(property -> configuration.set(Objects.toString(property.getKey()),
                Objects.toString(property.getValue())));
        this.connection = ConnectionFactory.createConnection(configuration);
        this.queue = new LinkedBlockingQueue<>();
        this.executor = Executors.newSingleThreadExecutor();
    }

    @Override
    public void start() throws Exception {
        isRunning = true;
        executor.execute(() -> {
            try (Table table = connection.getTable(TableName.valueOf(tableName.getBytes()))) {
                while (isRunning()) {
                    T t = queue.take();
                    Put put = new Put(keySupplier.apply(t));
                    put.addColumn(columnFamily.getBytes(), keySupplier.apply(t), System.currentTimeMillis(), valueSupplier.apply(t));
                    table.put(put);
                }
            } catch (InterruptedException | IOException e) {
                getLogger().error(e.getMessage(), e);
            }
        });
    }

    @Override
    public void close() throws Exception {
        try {
            isRunning = false;
            while (isRunning()) {
                TimeUnit.MILLISECONDS.sleep(500L);
            }
            executor.awaitTermination(500L, TimeUnit.MILLISECONDS);
        } finally {
            if (connection != null) {
                connection.close();
            }
        }
    }

    @Override
    public boolean isRunning() {
        return isRunning || !queue.isEmpty();
    }

    @Override
    public void put(T t) throws Exception {
        if (!isRunning()) {
            throw new IllegalStateException("Sink has been closed. ");
        }
        queue.offer(t);
    }

    public String getTableName() {
        return tableName;
    }

    public void setTableName(String tableName) {
        this.tableName = tableName;
    }

    public String getColumnFamily() {
        return columnFamily;
    }

    public void setColumnFamily(String columnFamily) {
        this.columnFamily = columnFamily;
    }

    public Function<T, byte[]> getKeySupplier() {
        return keySupplier;
    }

    public void setKeySupplier(Function<T, byte[]> keySupplier) {
        this.keySupplier = keySupplier;
    }

    public Function<T, byte[]> getValueSupplier() {
        return valueSupplier;
    }

    public void setValueSupplier(Function<T, byte[]> valueSupplier) {
        this.valueSupplier = valueSupplier;
    }
}
